/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.server;

import static org.sss.common.Utils.closeDataReceivers;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.UUID;

import org.slf4j.Logger;
import org.sss.common.Bucket;
import org.sss.common.CanceledException;
import org.sss.common.CommonStatus;
import org.sss.common.DataReceiver;
import org.sss.common.Inlet;
import org.sss.common.KeyValues;
import org.sss.common.KeyValuesPair;
import org.sss.common.TaskSliceSubmitter;
import org.sss.mapreduce.Configuration;
import org.sss.mapreduce.Decoder;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.Reducer;
import org.sss.mapreduce.SssException;
import org.sss.mapreduce.StorageNode;
import org.sss.mapreduce.datatype.Packable;
import org.sss.util.BytesComparator;
import org.sss.util.CallbackInterface;

public class MergeReduceTask extends ReduceTaskBase<KeyValuesPair> {
  private static Logger logger = SssServer.getLogger(MergeReduceTask.class);
  private static BytesComparator CMP = new BytesComparator();
  private final GroupID rightInputGID;
  private final ThreadLocal<Decoder> rightValueDecoder;
  private boolean handleOneSideOnly = false;
  private List<MergeSubmitter> mergeSubmitters = new ArrayList<MergeSubmitter>();

  public MergeReduceTask(TaskSet taskSet,
                         UUID jobID,
                         SssServerResource resource,
                         Configuration conf,
                         StorageNode storageNode,
                         Inlet leftInlet,
                         Inlet rightInlet,
                         Class<? extends Reducer> reducerClass,
                         List<org.sss.common.Outlet> outletList,
                         CallbackInterface<CommonStatus> doneCallback) throws SssException {
    super(taskSet, jobID, resource, conf, storageNode, 
        leftInlet, reducerClass, outletList, doneCallback);
    this.rightInputGID = rightInlet.GID;
    this.rightValueDecoder = loadThreadLocalDecoder(conf, rightInlet.valueClassName);
    this.handleOneSideOnly = conf.getBoolean("marge_reducer.handle_one_side_only", false);
    logger.info("handleOneSideOnly: " + handleOneSideOnly);
  }

  @Override
  public synchronized void renew() throws SssException {
    super.renew();
    mergeSubmitters.clear();
  }

  @Override
  protected void submitReadTaskSlice(int i) throws SssException {
    MergeSubmitter ms = new MergeSubmitter();
    mergeSubmitters.add(ms);
    submitReadRequest(getIOManager().createKeyValueReader(
        getConfiguration(), getStorageNode(), getInputGID(), i, ms.getLeft()));
    submitReadRequest(getIOManager().createKeyValueReader(
        getConfiguration(), getStorageNode(), rightInputGID, i, ms.getRight()));
  }

  @Override
  public void handle(List<KeyValuesPair> kvsps) {
    try {
      SssServerConfiguration sssConf = getSssServerConfiguration();
      List<DataReceiver<Bucket>> cache = getOutlets().getTLPutters(sssConf.reduce_output_cache_nbytes);
      Object[] args = createArguments(createContext(), 4, cache);

      final Decoder keyDecoder      = getKeyDecoder();
      final Decoder leftValDecoder  = getValueDecoder();
      final Decoder rightValDecoder = rightValueDecoder.get();

      Method reduceMethod = getProcInfo().getMethod();
      Reducer reducer = getProc();
      for (KeyValuesPair kvsp: kvsps) {
        Packable key = keyDecoder.decode(kvsp.key);
        args[1] = key;
        args[2] = decodingIterator(leftValDecoder, kvsp.leftValues);
        args[3] = key;
        args[4] = decodingIterator(rightValDecoder, kvsp.rightValues);
        checkCanceled();
        reduceMethod.invoke(reducer, args);
        kvsp.close();
      }
      closeDataReceivers(cache);
    } catch (Throwable e) {
      handleExceptionInHandle(e, "Error in merge reduce task. Cancels tasks.");
    } finally {
      for (KeyValuesPair kvsp: kvsps) {
        kvsp.close();
      }
      logger.debug("completing a ReduceTaskSlice.");
    }
  }

  @Override
  public synchronized void cancel() {
    super.cancel();
    for (MergeSubmitter merger: mergeSubmitters) {
      merger.cancel();
    }
  }

  private class MergeSubmitter {
    private TaskSliceSubmitter<List<KeyValuesPair>> submitter;
    private FragmentSubmitter left;
    private FragmentSubmitter right;
    private boolean isCanceled = false;

    public MergeSubmitter() {
      this.submitter = new TaskSliceSubmitter<List<KeyValuesPair>>(getWKTP(), MergeReduceTask.this);
      this.left      = new FragmentSubmitter();
      this.right     = new FragmentSubmitter();
    }

    public synchronized void cancel() {
      if (!isCanceled) {
        isCanceled = true;
        notifyAll();
      }
    }

    private synchronized void merge() throws SssException {
      List<KeyValuesPair> pairs = new ArrayList<KeyValuesPair>();
      KeyValues lkvs = null;
      KeyValues rkvs = null;
      while (true) {
        lkvs = left.getCurrentKVS();
        rkvs = right.getCurrentKVS();
        if (lkvs == null || rkvs == null) {
          break;
        }
        int order = CMP.compare(lkvs.key, rkvs.key);
        if (order == 0) {
          addPair(pairs, left.pop(), right.pop());
        }
        else if (order < 0) {
          addPair(pairs, left.pop(), null);
        }
        else { // order > 0
          addPair(pairs, null, right.pop());
        }
      }
      if (rkvs == null && right.closed) {
        while ((lkvs = left.getCurrentKVS()) != null) {
          addPair(pairs, left.pop(), null);
        }
      }
      if (lkvs == null && left.closed) {
        while ((rkvs = right.getCurrentKVS()) != null) {
          addPair(pairs, null, right.pop());
        }
      }

      if (!pairs.isEmpty()) {
        submitter.put(pairs);
      }
      if (lkvs == null || rkvs == null) {
        notifyAll();
      }
      if (left.closed && right.closed) {
        submitter.close();
      }
    }

    /**
     * .
     * WARNING: it is necessary to have lock.
     */
    private void addPair(List<KeyValuesPair> pairs, KeyValues left, KeyValues right) {
      assert left != null || right != null;
      if (handleOneSideOnly || (left != null && right != null)) {
        pairs.add(new KeyValuesPair(
              (left != null)?left.key:right.key,
              (left != null)?left.values:null,
              (right != null)?right.values:null));
        left = null;
        right = null;
      }
      if (left != null) left.values.close();
      if (right != null) right.values.close();
    }

    private synchronized void putToFragment(FragmentSubmitter f, List<KeyValues> kvs) throws SssException {
      try {
        while (f.tuples != null && !isCanceled) {
          wait();
        }
        if (isCanceled) {
          throw new CanceledException("MergeSubmitter is canceled");
        }
      }
      catch (InterruptedException e) {
        logger.error("Unexpected InterruptedException.");
        for (KeyValues kv: kvs) {
          kv.values.close();
        }
        submitter.setError(e);
        return;
      }
      f.tuples = kvs.iterator();
      assert f.tuples.hasNext(): "kvs has no items";
      f.cur = f.tuples.next();
      merge();
    }

    public DataReceiver<List<KeyValues>> getLeft()  { return left; }
    public DataReceiver<List<KeyValues>> getRight() { return right;}

    private class FragmentSubmitter implements DataReceiver<List<KeyValues>> {
      private KeyValues cur = null;
      private Iterator<KeyValues> tuples = null;
      private boolean closed = false;

      @Override
      public void put(List<KeyValues> kvs) throws SssException {
        if (kvs.isEmpty()) {
          return;
        }
        putToFragment(this, kvs);
      }

      public KeyValues getCurrentKVS() {
        return cur;
      }

      /**
       * .
       * WARNING: it is necessary that MergeSubmitter is locked.
       */
      private KeyValues pop() {
        KeyValues old = cur;
        if (tuples != null && tuples.hasNext()) {
          cur = tuples.next();
        }
        else {
          cur = null;
          tuples = null;
        }
        return old;
      }

      @Override
      public void close() throws SssException {
        synchronized (MergeSubmitter.this) {
          if (!closed) {
            closed = true;
            merge();
          }
        }
      }

      @Override
      public void setError(Throwable e) {
        submitter.setError(e);
      }

      @Override
      public synchronized Throwable getError() {
        return submitter.getError();
      }
    }
  }
}
